% ns_time : time spend on each step   %%%%%%%%FIX%%%%%%%%%%%%
% ns_slope: hip angle related outputs
% hip_pos: stance ankle related outputs
% pos_com: torso related angles  %%%%%%%%FIX%%%%%%%%%%%%
% s_knee: stance knee angle
% ns_knee: non stance knee angle

% LineNumber: output combination number
function optimization_MEANHUMAN(LineNumber)
global ns_time ns_slope s_knee ns_knee hip_pos pos_com
opt_addpath()

% initial condtion
a_fit = set_a_fit();

%%%%%%%%%%%%%% Other conditions for the optimization:
% [ub,lb]  = set_ublb(0.8);
[ub,lb]  = set_ublb(0.7);
Aeq = [];
beq = [];
A = [];
b = [];
set_globalveriable(LineNumber);


%%%%% Conditions for optimization: tight error
errorbound = 1e-8;
options = optimset('Algorithm','sqp',...
    'Display','iter','TolX',errorbound,...
    'MaxFunEvals',20000,'TolCon',errorbound);


% optimization
[a_out,cost_out] = fmincon(@cost,a_fit,A,b,Aeq,beq,lb,ub,@mycon,options);


%%%%%%%%%%%%%% Output the parameter matrix and the fit matrix:
a_opt = unflatten(a_out);
a_fit = unflatten(a_fit);
%%%%%%%%%%%%%% Output the initial condition:
[x_opt,tau_opt] = hzd_con(a_opt);
save('x_opt.mat','x_opt');
save('a_opt.mat','a_opt');

opt_plot(LineNumber,a_opt,a_fit);
%%%%%%%%%%%%%% Run main
OutPutType = num2str(LineNumber);
main_NAO(OutPutType)
end
%% cost function
function ret = cost(aflat)
global ns_time s_time ns_slope s_knee ns_knee hip_pos pos_com
% global kneemin hmin hmax qdotmax umax deltaP powermax kneemax

%%%%%%%%%%%%%% Unflatten the a matrix
a = unflatten(aflat);

%%%%%Pick the cost function
%%% 0 = torque based cost, requires additional integration
%%% 1 = human data based cost
typeofcost = 1;

%%%%%The torque based cost
if typeofcost == 0
    
    ret = torque_cost(a);
    
    %%%% The data based cost
elseif typeofcost == 1
    
    ret = humandata_cost(a);
end

end

function ret = torque_cost(a)
%%% Determine the point on the guard through inverse kinematics
[x_minus,Ttau,yminus,H,phipcond] = hzd_con(a);
%%% Integrate the system forward one step
x_plus = resetFunc(x_minus);
[Tf,xf] = step_under_reduced(x_plus,a,Ttau);

%%% Determine the step length (if needed
posf = jpos_mat(xf);
steplength = posf(1,end);
if steplength < 0
    steplength = 0.0001;
end

%%% Compute the cost
% ret =  norm(xf-x_minus)^2;
% ret = umax/abs(steplength)+norm(xf-x_minus)^2;
% ret = umax/abs(steplength);
%  ret = umax/steplength;
ret = (umax)/steplength;
end

function ret = humandata_cost(a)
global ns_time s_time ns_slope s_knee ns_knee hip_pos pos_com
%%%% Decide if you want to scale time based upon the parameterization of
%%%% time
scaletime = 0;
if scaletime == 0
    
    time_scale = ns_time;
    
elseif scaletime == 1

    [x_minus,Ttau,yminus,H,phipcond] = hzd_con(a);
    time_scale = ns_time*(Ttau/ns_time(end));
    
end
%     size(ns_slope)
%     size(ns_time)
%%%%%%%%%%%%%% Cost for each canonical function

cost_ns_slope = sum(abs(ns_slope-yd_ns_slope(time_scale,a)));
cost_hip_pos = sum(abs(hip_pos-yd_hip_pos(time_scale,a)));
cost_s_knee = sum(abs(s_knee-yd_s_knee(time_scale,a)));
cost_ns_knee = sum(abs(ns_knee-yd_ns_knee(time_scale,a)));
cost_pos_com = sum(abs(pos_com-yd_pos_com(time_scale,a))); %1007
% % cost_pos_com = sum(abs(0.02 - yd_pos_com(time_scale,a)));

%%%%%%%%%%%%%% Wieghting for the fits (in this case the torso cost is
%%%%%%%%%%%%%% not used.

beta_ns_slope = 1/(max(ns_slope)-min(ns_slope));
beta_hip_pos = 1/(max(hip_pos)-min(hip_pos));
beta_s_knee = 1/(max(s_knee)-min(s_knee));
beta_ns_knee = 1/(max(ns_knee)-min(ns_knee));
beta = beta_ns_slope+beta_hip_pos+beta_s_knee+beta_ns_knee;
% % % %         beta_pos_com = 1; %/(max(pos_com)-min(pos_com));
beta_pos_com = 1/(max(pos_com)-min(pos_com));
beta = beta_ns_slope+beta_hip_pos+beta_s_knee+beta_ns_knee+beta_pos_com;

%%% Determine the cost
cost_all = [
    beta_ns_slope*cost_ns_slope,...
    beta_hip_pos*cost_hip_pos,...
    beta_s_knee*cost_s_knee,...
    beta_ns_knee*cost_ns_knee,...
    beta_pos_com*cost_pos_com
    ];

%%%%%%%%%%%%%% Return the cost
ret = sum(cost_all)/beta;
end

%%
function [c,ceq] = mycon(aflat)
global ns_time 

%%%%%%%%%%%%%% Unflatten the a matrix
a = unflatten(aflat);

%%
%%%%Read in the constraints for HZD
[x_minus,Ttau,yminus,H] = hzd_con(a);
x_plus = resetFunc(x_minus);

%%
%%%%%%%%%%%%%% Require the post impact velocities be completlty determined:
%%% partial hybrid zero dynamics
% Hzd = H(2:end,:) - [0; 0; 0; 0];
%%% full hybrid zero dynamics
Hzd = H - [a(1,1); 0; 0; 0; 0];


%%%%%%%%%%%%%%%%%%%%% height %%%%%%%%%%%%%%%%%%%%%%%%%
 hmax = 0.02;
 ttime = [];
 h = [];
 hquad = [];
 hcond = [];
 
 
for t = 0:0.01:Ttau+0.01
    ttime = [ttime, t];
    h = [h, height(t,a,p)];
    hquad = [hquad, quad_height(t,hmax,Ttau)];
    hcond = [hcond, height(t,a,p) - quad_height(t,hmax,Ttau)];
end

hmaxcond = max(h);
hcondmin = min(hcond);

%%%%%%%%%%%%%% Pre and Post Impact knee angles

%%%%%%%%%%%%%% Nonlinear inequalities
%%% This is where you specify any constraints on the system
dimq = length(x_minus)/2;

    %%%%%%%%%%%%%% Nonlinear inequalities
    %%%%%%%%%%%%%%%%% gaurd must be negative
    %%%%%%%%%%%%%%%%% Time must be less than the human time
    c = [
           h_dot_mat(x_minus(1:dimq))*x_minus(dimq+1:end)+.01;
           -hcondmin;
        ];
    
%%%%%%%%%%%%%% Nonlinear equalities
%%%%%%%%%%%%%%%%% y pre impact must be zero
%%%%%%%%%%%%%%%%% y dot post impact must be zero
%%%%%%%%%%%%%%%%% guard must be satisfied
ceq = [
       yminus;
       Ttau-ns_time(end);
%        a(1,2) - deltaphip_sca(x_plus);
       Hzd];

end


function h = quad_height(t,hmax,Ttau)

a = -4*hmax/(Ttau^2);
b = 4*hmax/Ttau;

h = a*t^2 + b*t;


end

function h = quad_height_foot(q,hmax,q0,qf)

a = -((4*hmax)/(q0 - qf)^2);
b = (4*hmax*(q0 + qf))/(q0 - qf)^2;
c = -((4*hmax*q0*qf)/(q0 - qf)^2);

h = a*q^2 + b*q + c;


end

function [x_minus,Ttau,yminus,H,phipcond] = hzd_con(a)

%%%%%%%%%%Determine the initial position on the guard

q_minus = inverse_kin(a);

%%
%%%%%%%%%%Determine the pre and post impact values of the output functions

%%%%Initial position of the hip
dimq = length(q_minus);
q_plus = R(q_minus);
phipcond = phip_sca(q_plus);

%%%%Pre-impact relative degree 2 outputs
yminus = ya2_vec(q_minus, phipcond, a)-yd2_vec(q_minus, phipcond,  a);
% y2plus = ya2_vec(q_plus, phipcond, a)-yd2_vec(q_plus, phipcond,  a)

%%%%Pre-and post impact relative Jacobians of the outputs

Dy2_plus = Dya2_mat(q_plus, phipcond, a)-Dyd2_mat(q_plus, phipcond,  a);
Hplus = [ phip_dot_mat(q_plus);
    Dy2_plus(:,1:dimq)];

Dy2_minus = Dya2_mat(q_minus, phipcond, a)-Dyd2_mat(q_minus, phipcond,  a);
Hminus = [ phip_dot_mat(q_minus);
    Dy2_minus(:,1:dimq)];

%%%%%%%%%%The Time to impact (assuming no scuffing)
Ttau = sigma_sca(q_minus,phipcond,a);

%%%%Pre- and Post-impact velocities

qdot_minus = Hminus\[a(1,1); 0; 0; 0; 0];
x_minus = [q_minus; qdot_minus];

H = Hplus*P(x_minus);

end

%% 
%%%%%%%Computing the position at impact through inverse kinematics

function ret = inverse_kin(a)

%%%%Choose the type of inverse kinematics, either numeric or closed form
%%% 1 = numeric solving
%%% 2 = closed form solution (currently only works for non-stance slope and
%%% linearized non-stance slope)
typeofsolve = 1;

if typeofsolve == 1 
    %%%%% Pick initial guess
    %%%%%q0 as computed from the human parameters
    ret = numeric_inverse_kin(a);
elseif typeofsolve == 2
    %%%%If using the old version (no torso) use this:
    % theta_a(a)
    %%%If using the new version with the torso, use this:
    %%%%(the calculation is broken up, saves a lot of computation time)    
%     q = zeros(5,1);
    ret = math_inverse_kin(a);
    
end
end
function ret = math_inverse_kin(a)
    q = zeros(5,1);
    q(2) = theta_a2(q,a);
    q(5) = theta_a5(q,a);
    q(1) = theta_a1(q,a);
    q(4) = theta_a4(q,a);
    q(3) = theta_a3(q,a);
    
    ret = q;
end
function ret = numeric_inverse_kin(a)
   q0 = q_initial();
   
    %%%Numerically solving for the initial condition
    
    options = optimset('Display','off');
    qzero = fsolve(@(q)inverse_kin_fun(q,a),q0,options);
    
    ret = qzero;
end


%%%%%%%%%% initial guess %%%%%%%%%%%%%%%%
function ret = q_initial()
    ret = [ -0.5138
        0.2776
        0.2870
        -0.4201
        0.1746];
end

function F = inverse_kin_fun(q,a)

phipcond = phip_sca(R(q));
y2plus = ya2_vec(R(q), phipcond, a)-yd2_vec(R(q), phipcond,  a);

F = [y2plus;
    h_sca(q)];

end
%%
%%%%%%%%%%%%%% Computing the position and velocity components of the impact
%%%%%%%%%%%%%% equations

function ret = R(q)

dimq = length(q);
x = [q; zeros(dimq,1)];
xplus = resetFunc(x);

ret = xplus(1:dimq);

end

function ret = P(x)

dimq = length(x)/2;
xplus = resetFunc(x);

ret = xplus(dimq+1:end);

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%% plot %%%%%%%%%%%%%%%%%%
function opt_plot(LineNumber,a_opt,a_fit)
global ns_time ns_slope s_knee ns_knee hip_pos pos_com
OutPutType = num2str(LineNumber);
fig12=figure(12); clf;
subplot(2,3,1); plot(ns_time,yd_hip_pos(ns_time,a_opt),'b',ns_time,yd_hip_pos(ns_time,a_fit),'r',ns_time,hip_pos,'k.')
title('Hip Position')
xlabel('time (s)');
% ylabel('Desired vs. Fit vs. Data')
legend({'Desired','Fit','Data'});
subplot(2,3,2); plot(ns_time,yd_ns_slope(ns_time,a_opt),'b',ns_time,yd_ns_slope(ns_time,a_fit),'r',ns_time,ns_slope,'k.')
title('Non-Stance Slope')
xlabel('time (s)');
% ylabel('Desired vs. Fit vs. Data')
legend({'Desired','Fit','Data'});
subplot(2,3,4); plot(ns_time,yd_s_knee(ns_time,a_opt),'b',ns_time,yd_s_knee(ns_time,a_fit),'r',ns_time,s_knee,'k.')
title('Stance Knee')
xlabel('time (s)');
% ylabel('Desired vs. Fit vs. Data')
legend({'Desired','Fit','Data'});
subplot(2,3,5); plot(ns_time,yd_ns_knee(ns_time,a_opt),'b',ns_time,yd_ns_knee(ns_time,a_fit),'r',ns_time,ns_knee,'k.')
title('Non-Stance Knee')
xlabel('time (s)');
% ylabel('Desired vs. Fit vs. Data')
legend({'Desired','Fit','Data'});
subplot(2,3,3); plot(ns_time,yd_pos_com(ns_time,a_opt),'b',ns_time,yd_pos_com(ns_time,a_fit),'r',ns_time,pos_com,'k.')
title('Torso')
xlabel('time (s)');
% ylabel('Desired vs. Fit vs. Data')
legend({'Desired','Fit','Data'});
% figname = strcat('result/test.pdf');
figname = strcat('result/',OutPutType,'S.pdf');
% saveas(fig12,'LNST4.pdf')
saveas(fig12,figname)
end
%%
%%%%%%%%%%%%%% Canonical human functions

function ret = yd_ns_slope(t,a)

ret =exp(-a(2,4)*t).*(a(2,1)*cos(a(2,2)*t)+a(2,3)*sin(a(2,2)*t))+a(2,5);

end

function ret = yd_hip_pos(t,a)

ret = a(1,1)*t+a(1,2);

end

function ret = yd_s_knee(t,a)

ret = exp(-a(3,4)*t).*(a(3,1)*cos(a(3,2)*t)+a(3,3)*sin(a(3,2)*t))+a(3,5);

end

function ret = yd_ns_knee(t,a)

ret = exp(-a(4,4)*t).*(a(4,1)*cos(a(4,2)*t)+a(4,3)*sin(a(4,2)*t))+a(4,5);

end

function ret = yd_pos_com(t,a)

ret = exp(-a(5,4)*t).*(a(5,1)*cos(a(5,2)*t)+a(5,3)*sin(a(5,2)*t))+a(5,5);

end
%%
%%%%%%%%%%%%%%%%%%%% set the value for the gloabal variables
function set_globalveriable(LineNumber)
global ns_time ns_slope s_knee ns_knee hip_pos pos_com
opt_addpath()
[hipdatatype,nsslopedatatype,torsodatatype]  = load_Index(LineNumber);
% hip position
hip_all = load_hipdata(hipdatatype);
hip_pos = hip_all.hip;

% non-stance slope
ns_slope = load_nsslopedata(nsslopedatatype);
% stance knee
s_knee = load_skneedata();
% non-stance knee
ns_knee = load_nskneedata();
% torso
pos_com = load_torsodata(torsodatatype);
% time
ns_time = hip_all.ns_time; 
end

function a_fit = set_a_fit()
    opt_addpath();
    load('Afit.mat', 'a0');
    a_fit = flatten(a0);
end

% load human data of stance knee
function ret = load_skneedata()
    load('mean_data_sknee.mat');
    ret = s_knee;
end

% load human data of non-stance knee
function ret = load_nskneedata()
    load('mean_data_nsknee.mat');
    ret = ns_knee;
end

% load the output combination index
function [hipdatatype,nsslopedatatype,torsodatatype]  = load_Index(LineNumber)
    load('IndexAll.mat')
    hipdatatype = IndexAll(LineNumber,1);
    nsslopedatatype = IndexAll(LineNumber,2);
    torsodatatype = IndexAll(LineNumber,3);
end

%%%%%%%%%%%%%%% constraints %%%%%%%%%%%%%%%%%%%%%%%%
function [ub,lb]  = set_ublb(vel_limit)
ubmat = 50*ones(5,5);
lbmat = -50*ones(5,5);
%%upper and lower bounds on the hip velocity of the robot
ubmat(1,1) = vel_limit; %%% SHU
lbmat(1,1) = 0.1;

%%Add any constraints to the torso fit function here
% ubmat(5,:) = 0.1*ones(1,5);
% lbmat(5,:) = -0.0*ones(1,5);
%%%
ubmat(1,3) = 0; ubmat(1,4) = 0; ubmat(1,5) = 0;
lbmat(1,3) = 0; lbmat(1,4) = 0; lbmat(1,5) = 0;
% if hipdatatype ~= 4
    lbmat(1,2) = 0
    ubmat(1,2) = 0
% end
%%%
ub = flatten(ubmat);
lb = flatten(lbmat);
end
%%
%%%%%%%%%%%%%% Flatten and unflatten the a matrix

function aflat = flatten(a)
aflat = a(:);

end

function aunflat = unflatten(aflat)

qdim = 5;
aunflat = reshape(aflat,qdim,5);

end
